#include "common.h"
#include "llama.h"
#include <cmath>
#include <cstdio>
#include <string>
#include <vector>
#include <cmath>
#include <cfloat>
#include <fstream>
#include <iostream>
#include <string.h>
#include <stdio.h>
#include <unistd.h>
using	namespace	std;
uint64_t	prng=time(NULL);
static inline uint64_t wyrand(uint64_t	*seed) {
    *seed+=0xa0761d6478bd642full;
    uint64_t  see1=*seed^0xe7037ed1a0b428dbull;
    see1*=(see1>>32)|(see1<<32);
    return	(*seed*((*seed>>32)|(*seed<<32)))^((see1>>32)|(see1<<32));
}
static inline double wy2u01(uint64_t r) {
    const double _wynorm=1.0/(1ull<<52);
    return (r>>12)*_wynorm;
}
struct	lmchat {
    const int n_len = 1024;	//	can be modified, large value costs more memory
    gpt_params params;
    llama_model_params model_params;
    llama_model *model;
    llama_context_params ctx_params;
    llama_context *ctx=NULL;
    llama_batch batch= {};
    vector<llama_token> his,query;
    vector<float>	vs;
    vector<bool>	mask;
    int	n_cur,eot;

    void	init(const	char	*M) {
        params.model=M;
        llama_backend_init();
        llama_numa_init(params.numa);
        model_params = llama_model_default_params();
        model_params.n_gpu_layers = 200;	//	For GPU usage
        model = llama_load_model_from_file(params.model.c_str(), model_params);
        ctx_params = llama_context_default_params();
        ctx_params.seed  = 1234;
        ctx_params.n_ctx = n_len;
        ctx_params.n_batch = n_len;
        ctx_params.n_threads = 4; //params.n_threads;
        ctx_params.n_threads_batch = 4; //params.n_threads_batch == -1 ? params.n_threads : params.n_threads_batch;
        ctx_params.type_k=GGML_TYPE_Q8_0;
        vs.resize(llama_n_vocab(model));
    }
    void	close(void) {
        llama_batch_free(batch);
        llama_free(ctx);
        llama_free_model(model);
        llama_backend_free();
        printf("\n");
    }

    void	reset(void) {
        if(ctx!=NULL)	llama_free(ctx);
        ctx = llama_new_context_with_model(model, ctx_params);
        if(batch.n_tokens)	llama_batch_free(batch);
        batch = llama_batch_init(n_len, 0, 1);
        mask.assign(llama_n_vocab(model),true);
        for(size_t	i=0;	i<mask.size();	i++) {
            string	s=llama_token_to_piece(ctx, i);
            for(size_t	j=0;	j<s.size();	j++)	if(s[j]>=0)	mask[i]=false;
            if(i==llama_token_eos(model))	mask[i]=true;
        }        
        his=::llama_tokenize(ctx, params.prompt, true);
        for (size_t i = 0; i < his.size()&&i<n_len; i++) llama_batch_add(batch, his[i], i, { 0 }, false);
        batch.logits[batch.n_tokens - 1] = true;
        llama_decode(ctx, batch);
        n_cur=batch.n_tokens;
    }

    llama_token	sample(bool	must) {
        auto	eos=llama_token_eos(model);
        if(n_cur>n_len)	return	llama_token_eos(model);
        memset(vs.data(),0,llama_n_vocab(model)*sizeof(llama_token));
        for(uint32_t	k=0;	k<his.size();	k++) {
            uint32_t	l=1;
            while(l<=k&&his[k-l]==his[his.size()-l])	l++;
            vs[his[k]]+=(expf(l)-1)/(his.size()-k);
        }
        auto   n_vocab = llama_n_vocab(model);
        auto * logits  = llama_get_logits_ith(ctx, batch.n_tokens - 1);
        float	ma=-FLT_MAX,sum=0,sum1=0;
        for (llama_token token_id = 0; token_id < n_vocab; token_id++) ma=fmaxf(ma,logits[token_id]=((must&&token_id==eos)||!mask[token_id])?-1000:(logits[token_id]-vs[token_id]));
        for (llama_token token_id = 0; token_id < n_vocab; token_id++)	sum+=(logits[token_id]=expf(logits[token_id]-ma));
        for (llama_token token_id = 0; token_id < n_vocab; token_id++)	sum1+=(logits[token_id]=fmaxf(logits[token_id]/sum-1.0f/n_vocab,0.0f));
        float	ran=wy2u01(wyrand(&prng))*sum1;
        sum=0;
        llama_token new_token_id=0;
        for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
            sum+=logits[token_id];
            if(sum>=ran) {
                new_token_id=token_id;
                break;
            }
        }
        return	new_token_id;
    }
    llama_token	decision(void) {
        auto	eos=llama_token_eos(model);
        if(n_cur>n_len)	return	llama_token_eos(model);
        memset(vs.data(),0,llama_n_vocab(model)*sizeof(llama_token));
        for(uint32_t	k=0;	k<his.size();	k++) {
            uint32_t	l=1;
            while(l<=k&&his[k-l]==his[his.size()-l])	l++;
            vs[his[k]]+=(expf(l)-1)/(his.size()-k);
        }
        auto   n_vocab = llama_n_vocab(model);
        auto * logits  = llama_get_logits_ith(ctx, batch.n_tokens - 1);
        float	ma=-FLT_MAX,sum=0,sum1=0;
        for (llama_token token_id = 0; token_id < n_vocab; token_id++){
	        string	s=llama_token_to_piece(ctx, token_id);
	        if(!(s.size()>=3&&(s.substr(0,3)=="正"||s.substr(0,3)=="反")))	logits[token_id]=-1000;
        	ma=fmaxf(ma,logits[token_id]);
        }	        
        for (llama_token token_id = 0; token_id < n_vocab; token_id++)	sum+=(logits[token_id]=expf(logits[token_id]-ma));
        for (llama_token token_id = 0; token_id < n_vocab; token_id++)	sum1+=(logits[token_id]=fmaxf(logits[token_id]/sum-1.0f/n_vocab,0.0f));
        float	ran=wy2u01(wyrand(&prng))*sum1;
        sum=0;
        llama_token new_token_id=0;
        for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
            sum+=logits[token_id];
            if(sum>=ran) {
                new_token_id=token_id;
                break;
            }
        }
        return	new_token_id;
    }    
    string	reply(string	s,	bool	dec=false) {
        query=::llama_tokenize(ctx, s, true);
        llama_batch_clear(batch);
        for (size_t i = 0; i < query.size()&&n_cur<n_len; i++) {
            llama_batch_add(batch, query[i], n_cur, { 0 }, false);
            his.push_back(query[i]);
            n_cur++;
        }
        if(batch.n_tokens)	batch.logits[batch.n_tokens - 1] = true;
        llama_decode(ctx, batch);
        
        llama_token new_token_id,	eos=llama_token_eos(model);
        string	ret;
        for(size_t	n=0;	n_cur<n_len;	n++) {
            new_token_id=n==0&&dec?decision():sample(n<1);
            if(new_token_id==eos)	break;
            llama_batch_clear(batch);
            llama_batch_add(batch, new_token_id, n_cur, { 0 }, true);
            his.push_back(new_token_id);
            n_cur++;
            llama_decode(ctx, batch);
            ret+=llama_token_to_piece(ctx, new_token_id);
        };
        return	ret;
    }
};

void	document(void) {
    cerr<<"usage:	debate [options] topic positive negative\n";
    cerr<<"\t-m:	model file=7B.gguf\n";
    exit(0);
}

int main(int ac, char ** av) {
    lmchat	model;
    int	opt;
    string	model_file="/data/7B.gguf";
    while((opt=getopt(ac,	av,	"m:"))>=0) {
        switch(opt) {
        case	'm':
            model_file=optarg;
            break;
        default:
            document();
        }
    }
    if(ac<optind+3) {
        document();
        return	0;
    }
    model.params.prompt="下面将进行一场辩论赛，要求双方紧密围绕主题，进行针锋相对的辩论，不允许媾和。\n辩论主题为：";
    model.params.prompt+=av[optind];
    model.params.prompt+="\n正方观点：";
    model.params.prompt+=av[optind+1];
    model.params.prompt+="\n反方观点：";
    model.params.prompt+=av[optind+2];
    model.params.prompt+="\n";
    
    model.init(model_file.c_str());
    model.reset();
    cout<<model.params.prompt<<'\n';
    for(size_t	it=0;	it<3;	it++){
		string	a=model.reply("\n\n正方：");
		cout<<"\n正方："<<a<<'\n';
		string	b=model.reply("\n\n反方：");
		cout<<"\n反方："<<b<<'\n';
		
	}
	string	r=model.reply("\n\n请评委根据双方辩论内容，在正反方中二选一，判出胜负。\n评委：获胜方为",true);
	cout<<"\n\n请评委根据双方辩论内容，在正反方中二选一，判出胜负。\n评委：获胜方为"<<r<<'\n';
	ofstream	fo("result.txt",ios::app);
	fo<<r.substr(0,3)<<'\n';
	fo.close();
    model.close();
    return	0;
}
